Flash attn backend seam - #16
Open
eric-tc-wong wants to merge 3 commits into
Open
Conversation
The flash_attn dependency was not abstracted: 7 modules imported it directly and 20 call sites each re-implemented the same three compatibility branches inline (the 8-vs-4 forward tuple unpack, the window_size arity split, and the torch.ops.flash_attn version gate). Introduce ring_flash_attn/flash_attn_backend.py, which normalises the four primitives the ring variants actually need -- fa_forward / fa_backward and their varlen flavours -- onto the FA2 keyword set, and route every call site through it. Backend selection is RING_FLASH_ATTN_BACKEND, defaulting to fa2, so existing behaviour is unchanged unless a caller opts in. FA4 notes: its cute interface returns a 2-tuple, needs an explicit return_lse=True (autograd.Function.forward runs with grad disabled, so the requires_grad shortcut never fires), takes no dropout/ALiBi, maps the -1 window sentinel to None, and registers no torch.library ops -- so it is not traceable under torch.compile. Its backward additionally requires sm90+. The loader raises rather than falling back between backends: flash-attn-4 installs into the flash_attn namespace, so a broken flash-attn 2.x makes FA4 unimportable too, and a silent fallback would change numerics and hardware requirements underneath the caller. Variants outside FA4 scope are refused via check_variant_supported rather than run with unverified numerics. ring_flash_attn_backward is in the validated set because both llama hybrid entry points route into it. Verified against FA2: CPU tests plus 11 GPU tests and 7 compile-leg runs at nproc_per_node 1, all metrics bit-identical to a stashed baseline except dq diff, which is nondeterministic under FA2's atomic dq accumulation (shown via repeat runs on both sides, and by confirming the seam's only added kwargs are parameters passed at their own declared defaults). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
FA3 ships as a separate flash_attn_3 package with its own interface. Add it as a
third backend behind the existing seam, selected with RING_FLASH_ATTN_BACKEND=fa3.
FA2 remains the default and is unchanged.
FA3 is a much closer fit for ring attention than FA4:
* LSE layout matches FA2 exactly -- (batch, nheads, seqlen) dense and
(nheads, total_q) varlen -- so update_out_and_lse needs no transpose
* dq/dk/dv are written in place (mutates_args), FA3's only mode, so there is
no copy-back shim
* the -1 window sentinel is native; mapping it to None as FA4 needs would
fail FA3's torch.library schema, which types these as plain int
* both ops are real torch.library custom ops with register_fake, so unlike
FA4 they are traceable under torch.compile
Three ways the signatures differ from FA2, each of which fails quietly rather
than loudly if copied:
* backward renames causal to is_causal
* backward puts cu_seqlens_*/max_seqlen_* BEFORE dq/dk/dv, where FA2 puts
dq/dk/dv straight after softmax_lse, so a positional call misbinds
* dq/dk/dv are Optional and the C++ allocates throwaways when omitted,
computing the gradient and discarding it with no error
The adapter therefore calls keyword-only and always passes dq/dk/dv. Window
arity and the causal kwarg name are resolved by introspecting the real
signature via get_default_args, whose _init_fn fallback already handles
CustomOpDef -- older FA3 releases (still present as vendored hopper/ trees in
several envs here) use plain functions, a window_size tuple, and causal.
Varlen is unified upstream: one op per direction, varlen selected by passing
cu_seqlens_q. The four seam entry points are kept, with the varlen pair as thin
wrappers.
is_fa3_available probes flash_attn_3._C rather than the interface module,
because a bare hopper/ source tree with no built extension otherwise looks
available and dies at the first kernel call.
_FA4_VALIDATED_VARIANTS becomes a per-backend dict; fa3 shares fa4's llama-path
scope. _set_window moves to module scope, now shared by fa2 and fa3.
FA3 cannot run on this machine (no flash_attn_3._C, and stock builds gencode
sm90a), so it ships unrun. test_fa3_arg_binding covers what is checkable
without a GPU: it loads the seam under fa3 against a checkout's real ops with
the extension stubbed, and binds the adapter's own params dicts against the
true signatures, asserting on BoundArguments identity rather than just that
bind succeeded.
FA2 verified unchanged: CPU tests plus 11 GPU tests and 7 compile-leg runs.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The build-flag guard was dead. It looked up FLASH_ATTENTION_DISABLE_*, the spelling of the env vars that *set* the flags, but the generated config keys are FLASHATTENTION_DISABLE_* with no underscore after FLASH (hopper/setup.py:91), so every lookup missed and no build was ever rejected. Values are real booleans (setup.py:51 computes them as `== "TRUE"`), so plain truthiness is right once the keys are. The guard also only looked for a top-level `flash_attn_config`, which is what hopper/setup.py's py_modules installs -- but the layout that exposes flash_attn_3.flash_attn_interface keeps its config alongside, so the guard was a no-op on exactly the install this backend targets. Derive the config module from whichever interface module resolved, then fall back to the top-level name, and treat any unrecognised shape as "no opinion" rather than a failed import. Tests cover all of it: the guard fires on DISABLE_VARLEN and DISABLE_BACKWARD, stays quiet on a healthy build, and tolerates a list/str/None/unknown-keys config. Also assert that every kwarg the adapter passes but does not explicitly set equals its own declared default -- the property that makes get_default_args' pass-everything approach a semantic no-op, previously verified for FA2 but not FA3. test.sh's cpu_tests ran `python -m unittest test.<name>`, which resolved to the CPython stdlib `test` package because there is no test/__init__.py here; all three entries errored on import and the block never tested anything. Run them as bare module names with test/ on the path, and add the fa3 test so it is not orphaned. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request refactors the backend integration for FlashAttention kernels, consolidating backend selection and simplifying the interface for calling FlashAttention operations. It replaces direct imports and keyword-argument construction with new wrapper functions (
fa_forward,fa_backward, etc.), and standardizes the way window size and other parameters are passed. The changes also add optional dependency documentation for alternative FlashAttention backends.Backend abstraction and refactoring:
Replaces direct imports of FlashAttention kernel functions with wrapper functions (
fa_forward,fa_backward,fa_varlen_forward,fa_varlen_backward) from the newflash_attn_backendmodule in all main attention modules (ring_flash_attn.py,ring_flash_attn_varlen.py,llama_fwd_ring_bwd_flash_attn.py,llama3_flash_attn_varlen.py). This centralizes backend selection and allows for easier swapping or patching of the backend. [1] [2] [3] [4]Removes manual construction of parameter dictionaries and output unpacking for FlashAttention calls, replacing them with direct positional and keyword argument calls to the new backend wrappers. This reduces boilerplate and makes the code easier to follow and less error-prone. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10]
Parameter handling improvements:
window_sizeparameter by passing it as a single argument instead of separate left/right values, simplifying the call interface throughout the codebase. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10]Backend support checks:
check_variant_supportedat the start of top-level forward and backward functions to ensure that the chosen backend supports the requested operation, improving robustness. [1] [2] [3]Optional dependency documentation:
pyproject.tomlto addfa3andfa4as optional dependencies, documenting how to install alternative FlashAttention backends for different hardware or versions.